use bytes::Bytes;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{SinkExt, StreamExt};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::time::sleep;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::Utf8Bytes;

type WebSocketStream =
    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;

#[async_trait::async_trait]
pub trait TextMessageHandler: Sync + Send {
    async fn handle(&self, message: Utf8Bytes);
}

pub struct WatchTask<T>
where
    T: TextMessageHandler,
{
    sender: Arc<Mutex<SplitSink<WebSocketStream, Message>>>,
    receiver: Arc<Mutex<SplitStream<WebSocketStream>>>,
    message_handler: Arc<T>,
}

impl<T> WatchTask<T>
where
    T: TextMessageHandler + 'static,
{
    pub fn new(websocket: WebSocketStream, message_handler: T) -> Self {
        let (sender, receiver) = websocket.split();
        Self {
            sender: Arc::new(Mutex::new(sender)),
            receiver: Arc::new(Mutex::new(receiver)),
            message_handler: Arc::new(message_handler),
        }
    }

    pub async fn run(&self) {
        let keepalive_task = KeepaliveTask::new(self.sender.clone());
        let message_handler_task = MessageHandlerTask::new(
            self.sender.clone(),
            self.receiver.clone(),
            self.message_handler.clone(),
        );

        let keepalive_handle = tokio::spawn(async move {
            keepalive_task.run().await;
        });

        let message_handler_handle = tokio::spawn(async move {
            message_handler_task.run().await;
        });

        tokio::select! {
            _ = keepalive_handle => {
                log::error!("保活异常结束");
            }
            _ = message_handler_handle => {
                 log::error!("消息处理异常结束");
            }
        }
    }
}

struct KeepaliveTask {
    sender: Arc<Mutex<SplitSink<WebSocketStream, Message>>>,
}

impl KeepaliveTask {
    pub fn new(sender: Arc<Mutex<SplitSink<WebSocketStream, Message>>>) -> Self {
        Self { sender }
    }

    pub async fn run(&self) {
        loop {
            log::debug!("发送ping");
            if let Err(e) = self
                .sender
                .lock()
                .await
                .send(Message::Ping(Bytes::from("ping")))
                .await
            {
                log::error!("发送ping失败： {:?}", e);
                return;
            }
            sleep(Duration::from_secs(15)).await;
        }
    }
}

struct MessageHandlerTask<T>
where
    T: TextMessageHandler,
{
    sender: Arc<Mutex<SplitSink<WebSocketStream, Message>>>,
    receiver: Arc<Mutex<SplitStream<WebSocketStream>>>,
    handler: Arc<T>,
}

impl<T> MessageHandlerTask<T>
where
    T: TextMessageHandler,
{
    pub fn new(
        sender: Arc<Mutex<SplitSink<WebSocketStream, Message>>>,
        receiver: Arc<Mutex<SplitStream<WebSocketStream>>>,
        handler: Arc<T>,
    ) -> Self {
        Self {
            sender,
            receiver,
            handler,
        }
    }

    async fn handle_message(&self, message: Message) {
        match message {
            Message::Text(message) => {
                self.handler.handle(message).await;
            }
            Message::Ping(_) => {
                log::debug!("接收到ping：{}", message);
                if let Err(e) = self
                    .sender
                    .lock()
                    .await
                    .send(Message::Pong(Bytes::from("pong")))
                    .await
                {
                    log::error!("发送pong失败：{}", e);
                }
            }
            Message::Pong(_) => {
                log::debug!("接收到pong：{}", message);
            }
            _ => {
                log::debug!("其他消息：{:?}", message);
            }
        }
    }

    pub async fn run(&self) {
        while let Some(message) = self.receiver.lock().await.next().await {
            match message {
                Ok(message) => {
                    self.handle_message(message).await;
                }
                Err(e) => {
                    log::error!("接收消息异常退出：{:?}", e);
                    return;
                }
            }
        }
    }
}
